import pandas as pd
import requests
import re
import time
from tqdm import tqdm  # For progress tracking

def search_pubchem(compound_name):
    """
    Search PubChem for a compound and return its details if found
    """
    # Search for the compound
    search_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{compound_name}/json"
    try:
        search_response = requests.get(search_url)
        if search_response.status_code != 200:
            return None
        
        # Get CID from search results
        data = search_response.json()
        cid = str(data['PC_Compounds'][0]['id']['id']['cid'])
        
        # Get compound details using CID
        property_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/property/MolecularFormula,CanonicalSMILES/JSON"
        prop_response = requests.get(property_url)
        if prop_response.status_code != 200:
            return None
        
        prop_data = prop_response.json()['PropertyTable']['Properties'][0]
        
        # Get CAS number from compound details
        synonyms_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/synonyms/JSON"
        syn_response = requests.get(synonyms_url)
        cas = None
        
        if syn_response.status_code == 200:
            synonyms = syn_response.json()['InformationList']['Information'][0]['Synonym']
            # Look for CAS number pattern (numbers-numbers-numbers)
            cas_pattern = re.compile(r'\d+-\d+-\d+')
            for synonym in synonyms:
                if cas_pattern.match(synonym):
                    cas = synonym
                    break
        
        return {
            'CID': cid,
            'CAS': cas,
            'MolecularFormula': prop_data.get('MolecularFormula'),
            'SMILES': prop_data.get('CanonicalSMILES')
        }
        
    except Exception as e:
        print(f"Error processing {compound_name}: {str(e)}")
        return None

def main():
    # Read the Excel file
    df = pd.read_excel('compound_names.xlsx')
    
    # Create empty lists to store results
    results = []
    
    # Process each compound with a progress bar
    for compound in tqdm(df['Compound Name']):
        # Add delay to respect PubChem's API rate limits
        time.sleep(0.1)  
        
        result = search_pubchem(compound)
        if result:
            result['Compound Name'] = compound
            results.append(result)
        else:
            results.append({
                'Compound Name': compound,
                'CID': None,
                'CAS': None,
                'MolecularFormula': None,
                'SMILES': None
            })
    
    # Create result DataFrame
    result_df = pd.DataFrame(results)
    
    # Reorder columns to put Compound Name first
    columns = ['Compound Name', 'CID', 'CAS', 'MolecularFormula', 'SMILES']
    result_df = result_df[columns]
    
    # Save to Excel
    result_df.to_excel('pubchem_results.xlsx', index=False)
    print("Processing complete. Results saved to 'pubchem_results.xlsx'")

if __name__ == "__main__":
    main()